package biz.datalk.industrialland.common.encrypt.sm;

import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.util.encoders.Hex;

import javax.crypto.*;
import javax.crypto.spec.IvParameterSpec;
import javax.crypto.spec.SecretKeySpec;
import java.security.*;

/**
 * <a href="http://gmssl.org/docs/sm3.html">GmSSL SM4</a>
 *
 * @author tarofang@163.com
 * @date 2022-06-23
 */
public class SM4Util extends GMBaseUtil {
    /** 算法名称 */
    public static final String ALGORITHM_NAME = "SM4";
    /** 参数分别代表 算法名称/加密模式/数据填充方式 */
    public static final String ALGORITHM_NAME_ECB_PADDING = "SM4/ECB/PKCS5Padding";
    public static final String ALGORITHM_NAME_ECB_NO_PADDING = "SM4/ECB/NoPadding";
    public static final String ALGORITHM_NAME_CBC_PADDING = "SM4/CBC/PKCS5Padding";
    public static final String ALGORITHM_NAME_CBC_NO_PADDING = "SM4/CBC/NoPadding";
    /** 16字节  128位 */
    public static final int DEFAULT_KEY_SIZE = 128;

    public static byte[] generateKey() throws NoSuchAlgorithmException, NoSuchProviderException {
        return generateKey(DEFAULT_KEY_SIZE);
    }

    public static byte[] generateKey(int keySize) throws NoSuchAlgorithmException, NoSuchProviderException {
        KeyGenerator kg = KeyGenerator.getInstance(ALGORITHM_NAME, BouncyCastleProvider.PROVIDER_NAME);
        kg.init(keySize, new SecureRandom());
        return kg.generateKey().getEncoded();
    }

    public static byte[] encryptECBPadding(byte[] key, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
        return sm4(ALGORITHM_NAME_ECB_PADDING,  Cipher.ENCRYPT_MODE, key, null, data);
    }

    public static byte[] decryptECBPadding(byte[] key, byte[] data) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidAlgorithmParameterException {
        return sm4(ALGORITHM_NAME_ECB_PADDING,  Cipher.DECRYPT_MODE, key, null, data);
    }

    public static byte[] encryptECBNoPadding(byte[] key, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
        return sm4(ALGORITHM_NAME_ECB_NO_PADDING,  Cipher.ENCRYPT_MODE, key, null, data);
    }

    public static byte[] decryptECBNoPadding(byte[] key, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
        return sm4(ALGORITHM_NAME_ECB_NO_PADDING,  Cipher.DECRYPT_MODE, key, null, data);
    }

    public static byte[] encryptCBCPadding(byte[] key, byte[] iv, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
        return sm4(ALGORITHM_NAME_CBC_PADDING,  Cipher.ENCRYPT_MODE, key, iv, data);
    }

    public static byte[] decryptCBCPadding(byte[] key, byte[] iv, byte[] data) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidAlgorithmParameterException {
        return sm4(ALGORITHM_NAME_CBC_PADDING,  Cipher.DECRYPT_MODE, key, iv, data);
    }

    public static byte[] encryptCBCNoPadding(byte[] key, byte[] iv, byte[] data) throws InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, IllegalBlockSizeException, BadPaddingException, InvalidAlgorithmParameterException {
        return sm4(ALGORITHM_NAME_CBC_NO_PADDING,  Cipher.ENCRYPT_MODE, key, iv, data);
    }

    public static byte[] decryptCBCNoPadding(byte[] key, byte[] iv, byte[] data) throws IllegalBlockSizeException, BadPaddingException, InvalidKeyException, NoSuchAlgorithmException, NoSuchProviderException, NoSuchPaddingException, InvalidAlgorithmParameterException {
        return sm4(ALGORITHM_NAME_CBC_NO_PADDING,  Cipher.DECRYPT_MODE, key, iv, data);
    }

    /**
     * @param algorithm (*) 算法名称/加密模式/数据填充方式
     * @param mode      (*) 加解密模式 {@link Cipher#ENCRYPT_MODE} {@link Cipher#DECRYPT_MODE}
     * @param key       (*) 加解密的秘钥
     * @param iv        加解密向量
     * @param input     (*) 明文（加密模式）或密文（解密模式）
     *
     * @return 密文（加密模式）或明文（解密模式）
     *
     * @throws NoSuchPaddingException 未知的填充方式
     * @throws NoSuchAlgorithmException 未知的算法
     * @throws NoSuchProviderException 未知的算法提供者
     * @throws InvalidAlgorithmParameterException 无效的算法
     * @throws InvalidKeyException 无效的秘钥
     * @throws IllegalBlockSizeException 非法块异常
     * @throws BadPaddingException 填充异常
     * @author tarofang@163.com
     * @date 2022年06月27日
     */
    private static byte[] sm4(String algorithm, int mode, byte[] key, byte[] iv, byte[] input) throws NoSuchPaddingException, NoSuchAlgorithmException, NoSuchProviderException, InvalidAlgorithmParameterException, InvalidKeyException, IllegalBlockSizeException, BadPaddingException {
        if (algorithm == null || "".equals(algorithm.trim())) {
            throw new IllegalArgumentException("Algorithm is blank!");
        }
        if (key == null || key.length == 0) {
            throw new IllegalArgumentException("key is blank!");
        }
        if (mode != Cipher.ENCRYPT_MODE && mode != Cipher.DECRYPT_MODE) {
            throw new IllegalArgumentException(String.format("Invalid mode! [%s]", mode));
        }
        Key sm4Key = new SecretKeySpec(key, ALGORITHM_NAME);
        Cipher cipher = Cipher.getInstance(algorithm, BouncyCastleProvider.PROVIDER_NAME);
        if (iv != null) {
            IvParameterSpec ivParameterSpec = new IvParameterSpec(iv);
            cipher.init(mode, sm4Key, ivParameterSpec);
        } else {
            cipher.init(mode, sm4Key);
        }
        return cipher.doFinal(input);
    }

    // =================================================================================================================

    public static String genHexKey() {
        byte[] bytes;
        try {
            bytes = generateKey();
        } catch (Exception ex) {
            throw new RuntimeException(ex);
        }
        // return ByteUtils.toHexString(bytes);
        return Hex.toHexString(bytes);
    }
}
